package com.zb.zeus.gateway.limiter;

import com.alibaba.fastjson.JSON;
import com.zb.zeus.common.core.utils.StringUtils;
import com.zb.zeus.common.core.web.domain.AjaxResult;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.RequestRateLimiterGatewayFilterFactory;
import org.springframework.cloud.gateway.filter.ratelimit.KeyResolver;
import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter;
import org.springframework.cloud.gateway.route.Route;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;
import java.util.Map;
import java.util.Objects;

/**
 * @author fujianjian
 * @date 2023/4/11 16:59
 **/
@Slf4j
@Component
public class ZeusRequestRateLimiterGatewayFilterFactory extends RequestRateLimiterGatewayFilterFactory {

    private final RateLimiter defaultRateLimiter;
    private final KeyResolver defaultKeyResolver;

    public ZeusRequestRateLimiterGatewayFilterFactory(RateLimiter defaultRateLimiter,
                                                      KeyResolver hostAddrKeyResolver) {
        super(defaultRateLimiter, hostAddrKeyResolver);
        this.defaultRateLimiter = defaultRateLimiter;
        this.defaultKeyResolver = hostAddrKeyResolver;
    }

    @Override
    public GatewayFilter apply(Config config) {
        RateLimiter rateLimiter = getOrDefault(config.getRateLimiter(), defaultRateLimiter);
        KeyResolver keyResolver = getOrDefault(config.getKeyResolver(), defaultKeyResolver);

        return (exchage, chain) -> keyResolver.resolve(exchage).flatMap(key->{
            String routeId = config.getRouteId();
            if (StringUtils.isEmpty(routeId)) {
                Route route = exchage.getAttribute(ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR);
                routeId = route.getId();
            }
            final String finalRouteId = routeId;
            return rateLimiter.isAllowed(routeId, key).flatMap((response) -> {
                RateLimiter.Response tmpResp = (RateLimiter.Response) response;
                for (Map.Entry<String, String> header : tmpResp.getHeaders().entrySet()) {
                    exchage.getResponse().getHeaders().add(header.getKey(), header.getValue());
                }
                if (tmpResp.isAllowed()) {
                    return chain.filter(exchage);
                }
                log.warn("the request had limiter with route:{} rate limiter config", finalRouteId);
                ServerHttpResponse exchangeResp = exchage.getResponse();
                // 修改HTTPCODE为500
                exchangeResp.setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
                if (!exchangeResp.getHeaders().containsKey("Content-Type")) {
                    exchangeResp.getHeaders().add("Content-Type", "application/json");
                }

                DataBuffer buffer = exchangeResp.bufferFactory()
                        .wrap((JSON.toJSONString(AjaxResult.error(4029, "操作太过频繁")))
                                .getBytes(StandardCharsets.UTF_8));
                return exchangeResp.writeWith(Mono.just(buffer));
            });

        });
    }

    private <T> T getOrDefault(T configValue, T defaultValue) {
        return Objects.isNull(configValue) ? defaultValue : configValue;
    }
}
